#
#  Copyright 2022 The Open Islands Authors. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#
import os.path

from pyspark import SparkContext
from pyspark.sql import SparkSession

from pyoi.constant import StorageEngine
from pyoi.types import DataIO
from pyoi.types import DataSchema
from pyoi.util import UriUtils, log_utils
from pyoi.util.base_utils import json_dumps
from pyoi.util.base_utils import json_loads

LOGGER = log_utils.getLogger()


class DataSchemaUtils:
    def __init__(self, spark_session: SparkSession, spark_context: SparkContext):
        self._spark_session = spark_session
        self._spark_context = spark_context

    @classmethod
    def load_schema(cls, data_input: DataIO) -> DataSchema:
        if data_input.schema is not None:
            raise RuntimeError(f"schema is already in DataIO")
        if data_input.engine == StorageEngine.LOCALFS:
            if os.path.exists(UriUtils.trim_file_uri(data_input.schema_uri)):
                with open(UriUtils.trim_file_uri(data_input.schema_uri), "r") as fr:
                    schema = DataSchema.from_dict(json_loads(fr.read()))
                data_input.schema = schema
                LOGGER.info(f"load schema from {data_input.schema_uri}")
                return schema
            else:
                raise RuntimeError(f"schema uri {data_input.schema_uri} not exists")
        else:
            raise RuntimeError(f"{data_input.engine} is not supported")

    @classmethod
    def save_schema(cls, schema: DataSchema, data_output: DataIO):
        if data_output.engine == StorageEngine.LOCALFS:
            with open(UriUtils.trim_file_uri(data_output.schema_uri), "w") as fw:
                fw.write(json_dumps(schema))
        else:
            raise RuntimeError(f"{data_output.engine} is not supported")
